package cc.git.liuyan.customeraiagent.core.suppertool;

import cc.git.liuyan.customeraiagent.core.embeddingmodel.*;
import cc.git.liuyan.customeraiagent.core.vectordatabase.MysqlSchemaVectorData;
import cc.git.liuyan.customeraiagent.core.vectordatabase.VectorDataBase;
import lombok.extern.slf4j.Slf4j;

import java.math.BigDecimal;
import java.util.List;

//Text2Sql超级工具
@Slf4j
public class Text2SqlSupperTool {
    private static HybridEmbeddingModel innerHybridEmbeddingModel;
    private static VectorDataBase innerVectorDataBase;

    public static void init(HybridEmbeddingModel hybridEmbeddingModel, VectorDataBase vectorDataBase) {
        innerHybridEmbeddingModel = hybridEmbeddingModel;
        innerVectorDataBase = vectorDataBase;
    }

    private static String[] getSchemaAndTableNameFromDdl(String ddl) {
        int startIndex = ddl.indexOf("`");
        int endIndex = ddl.indexOf("(");
        return ddl.substring(startIndex, endIndex).trim().replaceAll("`", "").split("\\.");
    }

    /**
     * 训练Remark
     */
    public static void trainRemark(Long knowledgeId, Long remarkId, String remark) {
        EmbeddingModelInputData inputData = new EmbeddingModelInputData("remark", remark);
        EmbeddingModelOutputData remarkVector = innerHybridEmbeddingModel.embed(inputData).get(0);
        log.info("Text2SqlSupperTool.train Remark");
        innerVectorDataBase.createText2SqlVectorCollection(innerHybridEmbeddingModel.getDenseEmbeddingModel().dimension(), innerHybridEmbeddingModel.getSpaseEmbeddingModel().dimension());
        innerVectorDataBase.deleteText2SqlTrainRemark(knowledgeId, remarkId);
        MysqlSchemaVectorData mysqlSchemaVectorData = new MysqlSchemaVectorData();
        mysqlSchemaVectorData.setKnowledgeId(knowledgeId);
        mysqlSchemaVectorData.setType("remark");
        mysqlSchemaVectorData.setSchema("");
        mysqlSchemaVectorData.setTableName("");
        mysqlSchemaVectorData.setDdl("");
        mysqlSchemaVectorData.setRemarkId(remarkId);
        mysqlSchemaVectorData.setRemark(remarkVector.getSegmentContent());
        mysqlSchemaVectorData.setDenseVectors(remarkVector.getDenseVectors());
        mysqlSchemaVectorData.setSparseVectors(remarkVector.getSparseVectors());
        mysqlSchemaVectorData.setScore(1L);
        innerVectorDataBase.trainText2Sql(mysqlSchemaVectorData);
        log.info("Text2SqlSupperTool.train Remark end success");
    }

    /**
     * 训练Mysql
     */
    public static void trainMysql(Long knowledgeId, EmbeddingModelMysqlInputData inputData) {
        //ddl必带schema
        List<EmbeddingModelOutputData> allDdlVectors = innerHybridEmbeddingModel.embed(inputData);
        log.info("Text2SqlSupperTool.train DDL总数{}", allDdlVectors.size());
        innerVectorDataBase.createText2SqlVectorCollection(innerHybridEmbeddingModel.getDenseEmbeddingModel().dimension(), innerHybridEmbeddingModel.getSpaseEmbeddingModel().dimension());
        for (EmbeddingModelOutputData ddlVector : allDdlVectors) {
            String[] schemaAndTable = getSchemaAndTableNameFromDdl(ddlVector.getSegmentContent());
            String ddlSchema = schemaAndTable[0];
            String ddlTable = schemaAndTable[1];
            innerVectorDataBase.deleteText2SqlTrainDdl(knowledgeId, ddlSchema, ddlTable);
            MysqlSchemaVectorData mysqlSchemaVectorData = new MysqlSchemaVectorData();
            mysqlSchemaVectorData.setKnowledgeId(knowledgeId);
            mysqlSchemaVectorData.setType("ddl");
            mysqlSchemaVectorData.setSchema(ddlSchema);
            mysqlSchemaVectorData.setTableName(ddlTable);
            mysqlSchemaVectorData.setDdl(ddlVector.getSegmentContent());
            mysqlSchemaVectorData.setRemarkId(0L);
            mysqlSchemaVectorData.setRemark("");
            mysqlSchemaVectorData.setDenseVectors(ddlVector.getDenseVectors());
            mysqlSchemaVectorData.setSparseVectors(ddlVector.getSparseVectors());
            mysqlSchemaVectorData.setScore(1L);
            innerVectorDataBase.trainText2Sql(mysqlSchemaVectorData);
            log.info("Text2SqlSupperTool.train DDL schema:{}  table:{}", ddlSchema, ddlTable);
        }
        log.info("Text2SqlSupperTool.train DDL end success");
    }

    public static EmbeddingModelOutputData question2Vectors(String userQuestion) {
        return innerHybridEmbeddingModel.embed(new EmbeddingModelInputData("remark", userQuestion)).get(0);
    }

    public static List<String> relationSchemaSearch(List<BigDecimal> denseVectors, List<BigDecimal> sparseVectors) {
        return innerVectorDataBase.relationSchemaSearch(denseVectors, sparseVectors);
    }

    public static List<String> relationRemarkSearch(List<BigDecimal> denseVectors, List<BigDecimal> sparseVectors) {
        return innerVectorDataBase.relationRemarkSearch(denseVectors, sparseVectors);
    }
}
